/*
 * alternatepath.h
 *
 *  Created on: Mar 16, 2011
 *      Author: tqlong
 */

#ifndef ALTERNATEPATH_H_
#define ALTERNATEPATH_H_

#include <vector>
#include <queue>
#include <algorithm>

/** Any type that implements
 *    double get(i, j)
 *    int n1() const
 *    int n2() const
 *  is a cost matrix type
 */
template <typename DATASET>
class DistanceWeight
{
public:
  typedef DATASET                              dataset_type;
  typedef typename DATASET::vec_type           vec_type;
protected:
  dataset_type &d1_, &d2_;
  double count_;
public:
  DistanceWeight(dataset_type &d1, dataset_type &d2)
  : d1_(d1), d2_(d2) { resetCount(); }

  /** return the weight(i, j) which is the distance of d1[i] and d2[j] */
  double get(int i, int j) {
    count_ += 1;
    return arma::norm(d1_[i]-d2_[j], 2);
  }

  /** return number of points in each set */
  int n1() const { return d1_.n(); }
  int n2() const { return d2_.n(); }

  /** reset the number of accesses to zero */
  void resetCount() { count_ = 0; }
  double count() const { return count_; }
};

template <typename COSTMATRIX>
class AlternatePath
{
public:
  typedef COSTMATRIX                           costmatrix_type;
  typedef std::vector<double>                  dvec_type;
  typedef std::vector<int>                     ivec_type;
protected:
  costmatrix_type &C_;
  dvec_type u_, slack_;
  ivec_type next_, prev_, back_;
  double DOUBLE_INFINITY, TIGHT_TOLERANCE;
  double bigM_;
public:
  AlternatePath(costmatrix_type& C)
  : C_(C), u_(n()*2), slack_(n()), next_(n()), prev_(n()), back_(2*n()),
    DOUBLE_INFINITY(std::numeric_limits<double>::infinity()),
    TIGHT_TOLERANCE(1e-10)
  {}

  /** The alternating path algorithm, basically,
   *  maintains u(1..n1), v(1..n2) of dual variables
   *  such that u(i) + v(j) <= c(i, j) and finds matching
   *  edges among tight edges where u(i) + v(j) == c(i, j)
   */
  void optimize()
  {
    initializeDual();
    int count = 0;
    std::fill(next_.begin(), next_.end(), -1);
    std::fill(prev_.begin(), prev_.end(), -1);
    while (count < n()) {
      int v = findAlternatePath();
      if (v != -1) {
        flipPath(v);
        count++;
        std::cerr << "alternate path found " << count << "\n";
      }
      else {
//        std::cerr << "change dual " << count << "\n";
        changeDual();
      }
    }
  }

  int next(int i) const { assert(i < n()); return next_[i]; }
  int prev(int j) const { assert(j < n()); return prev_[j]; }
  int n() const { return std::max(C_.n1(), C_.n2()); }

protected:
  void initializeDual()
  {
    bigM_ = -DOUBLE_INFINITY;
    for (int i = 0; i < C_.n1(); i++) u_[i] = 0.0;
    for (int j = 0; j < C_.n2(); j++) {
      double min = DOUBLE_INFINITY;
      for (int i = 0; i < C_.n1(); i++) {
        double d = cost(i, j);
        if (d < min) min = d;
        if (d > bigM_) bigM_ = d;
      }
      u_[right(j)] = min;
    }
    bigM_ *= n();
  }

  int right(int j) const { return j+n(); }
  double cost(int i, int j)
  {
    if (i >= C_.n1() || j >= C_.n2()) return bigM_;
    else return C_.get(i, j);
  }
  bool tight(int i, int j)
  {
    double d = u_[i] + u_[right(j)] - cost(i, j);
    return fabs(d) < TIGHT_TOLERANCE;
  }

  int findAlternatePath()
  {
    std::fill(back_.begin(), back_.end(), -1);
    std::fill(slack_.begin(), slack_.end(), DOUBLE_INFINITY);
    std::queue<int> q;
    for (int i = 0; i < n(); i++)
      if (next_[i] == -1) {                      // push all un-assigned left vertices to queue
        q.push(i);
        back_[i] = -2;                           // mark them as visited
        updateSlack(i);
        break;
      }

    while (!q.empty()) {
      int u = q.front(); q.pop();                // pop the queue
      if (u < n()) {                             // if left vertex
        for (int v = 0; v < n(); v++)
          if (back_[right(v)] == -1 && tight(u, v)) {  // if v not yet in queue and (u,v) is a tight edge
            q.push(right(v));
            back_[right(v)] = u;
            if (prev_[v] == -1) return v;        // if v is unassigned, an alternate path is found
          }
      }
      else {                                     // if right vertex
        int v = prev_[u-n()];                    // add the corresponding match on the left to queue
        q.push(v);
        back_[v] = u;
        updateSlack(v);
      }
    }
    return -1;
  }

  void updateSlack(int i)
  {
    for (int j = 0; j < n(); j++) if (back_[right(j)] == -1) {
      double s = cost(i, j) - u_[i] - u_[right(j)];
      if (slack_[j] > s)
        slack_[j] = s;
    }
  }

  void flipPath(int v)
  {
    v = right(v);
    while (v >= 0) {
      int vv = v-n();
      int uu = back_[v];
      next_[uu] = vv;
      prev_[vv] = uu;
      v = back_[uu];
    }
  }

  void changeDual()
  {
    double delta = DOUBLE_INFINITY;
    for (int j = 0; j < n(); j++) if (back_[right(j)] == -1)
      if (slack_[j] < delta) delta = slack_[j];
    for (int i = 0; i < n(); i++) if (back_[i] != -1)
      u_[i] += delta;
    for (int j = 0; j < n(); j++) if (back_[right(j)] != -1)
      u_[right(j)] -= delta;
  }
};

#endif /* ALTERNATEPATH_H_ */
